#pragma once
#include <zMat.hpp>

namespace zzz{

template<int D>
struct TensorVoter
{
  TensorVoter():Position(0),Lambda(0),Directions(0),T(0){}

  //init from position
  //and init as a uniform ball tensor
  TensorVoter(const Vector<D, double> &pos)
    :Position(pos),Lambda(1),Directions(0),T(0)
  {
    //ball init
    for (int i=0; i<D; i++)
      Directions(i,i)=1;
  }

  //init from a vector
  //and init as a stick tensor
  //stick from pos point to normal, therefore normal-pos is normal
  TensorVoter(const Vector<D, double> &pos, Vector<D, double> normal)
    :Position(pos),Lambda(0),Directions(0),T(0)
  {
    //stick init
    normal-=pos;

    Lambda[0]=normal.Normalize();

    Vector<D, double> ori(0);
    ori[0]=1;
    Matrix<D,D,double> mat=GetRotationBetweenVectors(ori,normal);

    for (int i=0; i<D; i++)
    {
      Directions(i,i)=1;
      Directions.Row(i)=mat*Directions.Row(i);
    }
  }

  //init from Tensor
  //you can save and load T
  TensorVoter(const Vector<D, double> &pos, const Matrix<D, D, double> &t)
    :Position(pos),T(t)
  {
    //init from Tensor
    Decomp();
  }


  void Decomp()
  {
    zMatrix<double> A(Dress(T)),U,S,VT;
    SVD(A,U,S,VT);
    Dress<zColMajor,double,D>(Lambda)=S;
    Dress(Directions)=Trans(U);
  }

  void Encode()
  {
    Matrix<D, D, double> stick(GetStickVote());
    Matrix<D, D, double> plate(GetPlateVote());
    Matrix<D, D, double> ball(GetBallVote());
    T=stick+plate+ball;
  }

  //StickVote part
  Matrix<D, D, double> GetStickVote()
  {
    //TODO Optimize
    Matrix<D, D, double> mat;
    double lambda=Lambda[0]-Lambda[1];
    for (int i=0; i<D; i++) for (int j=0; j<D; j++)
      mat(i,j)=lambda*Directions(0,i)*Directions(0,j);
    return mat;
  }

  //PlateVote part
  Matrix<D, D, double> GetPlateVote()
  {
    //TODO Optimize
    Matrix<D, D, double> mat(0);
    Matrix<D, D, double> Pk=Matrix<D,1,double>(&(Directions(0,0)))*Matrix<1,D,double>(&(Directions(0,0)));
    for (int i=1; i<D-1; i++)
    {
      double lambda=Lambda[i]-Lambda[i+1];
      Pk+=Matrix<D,1,double>(&(Directions(i,0)))*Matrix<1,D,double>(&(Directions(i,0)));
      mat+=Pk*lambda;
    }
    return mat;
  }

  //BallVote part
  Matrix<D, D, double> GetBallVote()
  {
    Matrix<D, D, double> Pk=Matrix<D,1,double>(&(Directions(0,0)))*Matrix<1,D,double>(&(Directions(0,0)));
    for (int i=1; i<D; i++)
    {
      Pk+=Matrix<D,1,double>(&(Directions(i,0)))*Matrix<1,D,double>(&(Directions(i,0)));
    }
    Matrix<D, D, double> mat(Pk*Lambda[D-1]);
    return mat;
  }


  bool operator==(const TensorVoter<D>&other){return Position==other.Position;}

//data
  //voter 3d position
  Vector<D,double> Position;
  //eigen values
  Vector<D,double> Lambda;
  //eigen vectors
  Matrix<D,D,double> Directions;
  //T
  Matrix<D, D, double> T;


};
}
